from __future__ import annotations

from copy import deepcopy
from typing import Any

import pytest

from marimo._data.get_datasets import (
    _db_type_to_data_type,
    get_databases_from_duckdb,
    get_datasets_from_variables,
    get_duckdb_databases_agg_query,
    get_table_columns,
    has_updates_to_datasource,
)
from marimo._data.models import Database, DataTable, DataTableColumn, Schema
from marimo._dependencies.dependencies import DependencyManager
from marimo._types.ids import VariableName
from tests._data.mocks import create_dataframes

HAS_DEPS = DependencyManager.duckdb.has()


@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
def test_has_updates_to_datasource() -> None:
    assert has_updates_to_datasource("hello") is False
    assert has_updates_to_datasource("ATTACH 'marimo.db'") is True
    assert has_updates_to_datasource("DETACH marimo") is True
    assert has_updates_to_datasource("CREATE TABLE cars (name TEXT)") is True


sql_query = """
CREATE TABLE all_types (
    col_boolean BOOLEAN,
    col_tinyint TINYINT,
    col_smallint SMALLINT,
    col_integer INTEGER,
    col_bigint BIGINT,
    col_hugeint HUGEINT,
    col_utinyint UTINYINT,
    col_usmallint USMALLINT,
    col_uinteger UINTEGER,
    col_ubigint UBIGINT,
    col_float FLOAT,
    col_double DOUBLE,
    col_decimal DECIMAL(18,3),
    col_varchar VARCHAR,
    col_date DATE,
    col_time TIME,
    col_timetz TIMETZ,
    col_timestamp TIMESTAMP,
    col_interval INTERVAL,
    col_blob BLOB,
    col_bit BIT,
    col_uuid UUID,
    col_json JSON,
    col_real REAL,
    col_numeric NUMERIC(10,2),
    col_char CHAR(10),
    col_bpchar BPCHAR(10),
    col_text TEXT,
    col_string STRING,
    col_timestamp_ns TIMESTAMP_NS,
    col_timestamp_ms TIMESTAMP_MS,
    col_timestamptz TIMESTAMPTZ,
    col_datetime DATETIME,
    col_array INTEGER[],
    col_array_with_size INTEGER[10],
    col_struct STRUCT(id INTEGER, name VARCHAR),
    col_map MAP(INTEGER, VARCHAR),
    col_union UNION(int INTEGER, varchar VARCHAR),
    col_enum ENUM('a', 'b', 'c'),
    col_uhugeint UHUGEINT
);

CREATE SCHEMA s1;
CREATE SCHEMA s2;

CREATE TABLE s1.t (id INTEGER PRIMARY KEY, other_id INTEGER);
CREATE TABLE s2.t (id INTEGER PRIMARY KEY, j VARCHAR);
"""

all_types_tables = [
    DataTable(
        name="all_types",
        source_type="duckdb",
        source="memory",
        num_rows=None,
        num_columns=40,
        variable_name=None,
        columns=[
            DataTableColumn(
                name="col_boolean",
                type="boolean",
                external_type="BOOLEAN",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_tinyint",
                type="integer",
                external_type="TINYINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_smallint",
                type="integer",
                external_type="SMALLINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_integer",
                type="integer",
                external_type="INTEGER",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_bigint",
                type="integer",
                external_type="BIGINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_hugeint",
                type="integer",
                external_type="HUGEINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_utinyint",
                type="integer",
                external_type="UTINYINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_usmallint",
                type="integer",
                external_type="USMALLINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_uinteger",
                type="integer",
                external_type="UINTEGER",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_ubigint",
                type="integer",
                external_type="UBIGINT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_float",
                type="number",
                external_type="FLOAT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_double",
                type="number",
                external_type="DOUBLE",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_decimal",
                type="number",
                external_type="DECIMAL(18,3)",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_varchar",
                type="string",
                external_type="VARCHAR",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_date",
                type="date",
                external_type="DATE",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_time",
                type="time",
                external_type="TIME",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_timetz",
                type="time",
                external_type="TIME WITH TIME ZONE",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_timestamp",
                type="datetime",
                external_type="TIMESTAMP",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_interval",
                type="datetime",
                external_type="INTERVAL",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_blob",
                type="string",
                external_type="BLOB",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_bit",
                type="string",
                external_type="BIT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_uuid",
                type="string",
                external_type="UUID",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_json",
                type="unknown",
                external_type="JSON",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_real",
                type="number",
                external_type="FLOAT",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_numeric",
                type="number",
                external_type="DECIMAL(10,2)",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_char",
                type="string",
                external_type="VARCHAR",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_bpchar",
                type="string",
                external_type="VARCHAR",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_text",
                type="string",
                external_type="VARCHAR",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_string",
                type="string",
                external_type="VARCHAR",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_timestamp_ns",
                type="datetime",
                external_type="TIMESTAMP_NS",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_timestamp_ms",
                type="datetime",
                external_type="TIMESTAMP_MS",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_timestamptz",
                type="datetime",
                external_type="TIMESTAMP WITH TIME ZONE",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_datetime",
                type="datetime",
                external_type="TIMESTAMP",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_array",
                type="unknown",
                external_type="INTEGER[]",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_array_with_size",
                type="unknown",
                external_type="INTEGER[10]",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_struct",
                type="unknown",
                external_type='STRUCT(id INTEGER, "name" VARCHAR)',
                sample_values=[],
            ),
            DataTableColumn(
                name="col_map",
                type="unknown",
                external_type="MAP(INTEGER, VARCHAR)",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_union",
                type="unknown",
                external_type='UNION("int" INTEGER, "varchar" VARCHAR)',
                sample_values=[],
            ),
            DataTableColumn(
                name="col_enum",
                type="string",
                external_type="ENUM('a', 'b', 'c')",
                sample_values=[],
            ),
            DataTableColumn(
                name="col_uhugeint",
                type="integer",
                external_type="UHUGEINT",
                sample_values=[],
            ),
        ],
    )
]

s1_tables = [
    DataTable(
        name="t",
        source_type="duckdb",
        source="memory",
        num_rows=None,
        num_columns=2,
        variable_name=None,
        columns=[
            DataTableColumn(
                name="id",
                type="integer",
                external_type="INTEGER",
                sample_values=[],
            ),
            DataTableColumn(
                name="other_id",
                type="integer",
                external_type="INTEGER",
                sample_values=[],
            ),
        ],
    )
]

s2_tables = [
    DataTable(
        name="t",
        source_type="duckdb",
        source="memory",
        num_rows=None,
        num_columns=2,
        variable_name=None,
        columns=[
            DataTableColumn(
                name="id",
                type="integer",
                external_type="INTEGER",
                sample_values=[],
            ),
            DataTableColumn(
                name="j",
                type="string",
                external_type="VARCHAR",
                sample_values=[],
            ),
        ],
    )
]

cleanup_query = """
DROP TABLE all_types;
DROP SCHEMA s1 CASCADE;
DROP SCHEMA s2 CASCADE;
"""


@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
def test_get_databases() -> None:
    import duckdb

    duckdb.execute(sql_query)

    assert get_databases_from_duckdb(connection=None) == [
        Database(
            name="memory",
            dialect="duckdb",
            schemas=[
                Schema(name="main", tables=all_types_tables),
                Schema(name="s1", tables=s1_tables),
                Schema(name="s2", tables=s2_tables),
            ],
            engine=None,
        )
    ]

    duckdb.execute(cleanup_query)


@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
def test_get_databases_with_no_tables() -> None:
    import duckdb

    in_memory_database = Database(
        name="memory",
        dialect="duckdb",
        schemas=[],
        engine=None,
    )

    # No connection, return in-memory database
    connection = None
    assert get_databases_from_duckdb(connection=connection) == [
        in_memory_database
    ]
    assert get_duckdb_databases_agg_query(
        connection=connection, engine_name=None
    ) == [in_memory_database]

    # Custom connection with no tables
    connection = duckdb.connect(":memory:")
    assert get_databases_from_duckdb(connection=connection) == [
        in_memory_database
    ]
    assert get_duckdb_databases_agg_query(
        connection=connection, engine_name=None
    ) == [in_memory_database]


@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
def test_get_databases_with_connection() -> None:
    import duckdb

    connection = duckdb.connect(":memory:")
    connection.execute(sql_query)

    all_tables = deepcopy(all_types_tables)
    for table in all_tables:
        table.source_type = "connection"
        table.engine = VariableName("engine")

    s1 = deepcopy(s1_tables)
    for table in s1:
        table.source_type = "connection"
        table.engine = VariableName("engine")

    s2 = deepcopy(s2_tables)
    for table in s2:
        table.source_type = "connection"
        table.engine = VariableName("engine")

    expected_databases = [
        Database(
            name="memory",
            dialect="duckdb",
            schemas=[
                Schema(name="main", tables=all_tables),
                Schema(name="s1", tables=s1),
                Schema(name="s2", tables=s2),
            ],
            engine=VariableName("engine"),
        )
    ]

    assert (
        get_databases_from_duckdb(
            connection=connection, engine_name=VariableName("engine")
        )
        == expected_databases
    )

    assert (
        get_duckdb_databases_agg_query(
            connection=connection, engine_name=VariableName("engine")
        )
        == expected_databases
    )

    connection.execute(cleanup_query)


@pytest.mark.parametrize(
    "df",
    create_dataframes({"A": [1, 2, 3], "B": ["a", "a", "a"]}),
)
def test_get_datasets_from_variables(df: Any) -> None:
    datatests = get_datasets_from_variables(
        [(VariableName("my_df"), df), (VariableName("non_df"), 123)]
    )
    # We don't compare these values
    external_type1 = datatests[0].columns[0].external_type
    external_type2 = datatests[0].columns[1].external_type

    rows = datatests[0].num_rows
    assert rows is None or rows == 3

    sample_values1 = datatests[0].columns[0].sample_values
    sample_values2 = datatests[0].columns[1].sample_values
    assert sample_values1 == [1, 2, 3] or sample_values1 == []
    assert sample_values2 == ["a", "a", "a"] or sample_values2 == []

    assert datatests == [
        DataTable(
            name="my_df",
            source_type="local",
            source="memory",
            num_rows=rows,
            num_columns=2,
            variable_name=VariableName("my_df"),
            columns=[
                DataTableColumn(
                    name="A",
                    type="integer",
                    external_type=external_type1,
                    sample_values=sample_values1,
                ),
                DataTableColumn(
                    name="B",
                    type="string",
                    external_type=external_type2,
                    sample_values=sample_values2,
                ),
            ],
        )
    ]


def test_get_table_columns() -> None:
    import duckdb

    connection = duckdb.connect(":memory:")
    connection.execute(sql_query)

    columns = get_table_columns(connection, "all_types")
    assert columns == all_types_tables[0].columns


def test_get_databases_agg_query() -> None:
    import duckdb

    connection = duckdb.connect(":memory:")
    connection.execute(sql_query)

    assert get_duckdb_databases_agg_query(
        connection=connection, engine_name=None
    ) == [
        Database(
            name="memory",
            dialect="duckdb",
            schemas=[
                Schema(name="main", tables=all_types_tables),
                Schema(name="s1", tables=s1_tables),
                Schema(name="s2", tables=s2_tables),
            ],
            engine=None,
        )
    ]

    connection.execute(cleanup_query)


@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
def test_get_databases_with_closed_connection() -> None:
    """Test that closed connections are handled gracefully without errors."""
    import duckdb

    # Create a connection, add tables, then close it
    connection = duckdb.connect(":memory:")
    connection.execute("CREATE TABLE test_table (id INTEGER, name VARCHAR)")
    connection.close()

    # This should not raise an exception
    result = get_databases_from_duckdb(
        connection=connection, engine_name=VariableName("closed_engine")
    )

    # Should return empty list for closed connection
    assert result == []


@pytest.mark.skipif(not HAS_DEPS, reason="optional dependencies not installed")
def test_get_databases_with_context_manager_closed_connection() -> None:
    """Test that connections closed via context manager are handled gracefully."""
    import duckdb

    # Use context manager to automatically close connection
    with duckdb.connect(":memory:") as connection:
        connection.execute(
            "CREATE TABLE test_table (id INTEGER, name VARCHAR)"
        )
        # Store reference to connection
        closed_conn = connection

    # Connection is now closed, this should not raise an exception
    result = get_databases_from_duckdb(
        connection=closed_conn, engine_name=VariableName("closed_engine")
    )

    # Should return empty list for closed connection
    assert result == []


def test_db_type_to_data_type_null() -> None:
    """Test that the 'null' DuckDB type is handled without warnings."""
    # Test that null type maps to unknown
    assert _db_type_to_data_type("null") == "unknown"
    assert _db_type_to_data_type("NULL") == "unknown"
    assert _db_type_to_data_type('"null"') == "unknown"


def test_db_type_to_data_type_various() -> None:
    """Test various DuckDB type mappings."""
    # Integer types
    assert _db_type_to_data_type("integer") == "integer"
    assert _db_type_to_data_type("bigint") == "integer"
    assert _db_type_to_data_type("int128") == "integer"
    assert _db_type_to_data_type("integral") == "integer"
    assert _db_type_to_data_type("long") == "integer"
    assert _db_type_to_data_type("short") == "integer"
    assert _db_type_to_data_type("signed") == "integer"
    assert _db_type_to_data_type("oid") == "integer"
    assert _db_type_to_data_type("varint") == "integer"

    # Unsigned integers
    assert _db_type_to_data_type("utinyint") == "integer"
    assert _db_type_to_data_type("usmallint") == "integer"
    assert _db_type_to_data_type("uinteger") == "integer"
    assert _db_type_to_data_type("ubigint") == "integer"
    assert _db_type_to_data_type("uhugeint") == "integer"
    assert _db_type_to_data_type("uint128") == "integer"

    # Numeric types
    assert _db_type_to_data_type("float") == "number"
    assert _db_type_to_data_type("double") == "number"
    assert _db_type_to_data_type("float4") == "number"
    assert _db_type_to_data_type("dec") == "number"
    assert _db_type_to_data_type("decimal") == "number"

    # String types
    assert _db_type_to_data_type("varchar") == "string"
    assert _db_type_to_data_type("text") == "string"
    assert _db_type_to_data_type("blob") == "string"
    assert _db_type_to_data_type("guid") == "string"
    assert _db_type_to_data_type("nvarchar") == "string"

    # Binary types (represented as string)
    assert _db_type_to_data_type("binary") == "string"
    assert _db_type_to_data_type("varbinary") == "string"
    assert _db_type_to_data_type("bytea") == "string"

    # Boolean
    assert _db_type_to_data_type("boolean") == "boolean"
    assert _db_type_to_data_type("bool") == "boolean"
    assert _db_type_to_data_type("logical") == "boolean"

    # Date/Time
    assert _db_type_to_data_type("date") == "date"
    assert _db_type_to_data_type("timestamp") == "datetime"
    assert _db_type_to_data_type("timestamptz") == "datetime"
    assert _db_type_to_data_type("timetz") == "time"

    # Special types
    assert _db_type_to_data_type("geometry") == "unknown"
    assert _db_type_to_data_type("null") == "unknown"
    assert _db_type_to_data_type("json") == "unknown"
    assert _db_type_to_data_type("row") == "unknown"
